package cas.ihep.spark.loop;

import cas.ihep.spark.io.Packet;
import cas.ihep.spark.util.GlobalConfiguration;
import cas.ihep.spark.util.LogEnv;
import cas.ihep.spark.util.function.Function2;
import org.apache.spark.SparkEnv;
import org.apache.spark.executor.CoarseGrainedExecutorBackend;
import org.apache.spark.executor.Executor;
import org.apache.spark.rpc.netty.Dispatcher;
import org.apache.spark.rpc.netty.NettyRpcEnv;

import java.io.IOException;
import java.lang.reflect.InvocationTargetException;
import java.lang.reflect.Method;
import java.net.SocketAddress;
import java.util.HashMap;
import java.util.Map;
import java.util.Queue;
import java.util.concurrent.ArrayBlockingQueue;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.LinkedBlockingDeque;
import java.util.concurrent.atomic.AtomicInteger;

public class IterationClient<T,R> {
    private AtomicInteger pooledCount;
    private Map<Thread,BlockingQueue<IterationData<T,R>>> pooledThread;
    private Class<T> tClass;
    private SocketAddress localAddress,remoteAddress;

    private final static IterationData<?,?> EOFObject=new IterationData<Object, Object>() {
        @Override
        public Object input() {
            throw new RuntimeException("Should never call this method of an EOF object");
        }

        @Override
        public void output(Object ret) {
            throw new RuntimeException("Should never call this method of an EOF object");
        }
    };

    private IterationClient(){
        pooledCount=new AtomicInteger(0);
        pooledThread=new ConcurrentHashMap<>();
    }

    private static class  WrappedIterationData <T,R> implements IterationData<T,R>{
        T t;
        Queue<R> rValues;
        IterationNetClient.VirtualNetContext context;
        SocketAddress to;
        R zero;
        Function2<R,R,R> aggFunc;
        int numOfIters,pooledCount;

        WrappedIterationData(T tt, IterationNetClient.VirtualNetContext ctx,SocketAddress _2,R _0,Function2<R,R,R> agg,int iters,int pooled) {
            t=tt;
            rValues=new ArrayBlockingQueue<>(pooled);
            context=ctx;
            to=_2;
            zero=_0;
            aggFunc=agg;
            numOfIters=iters;
            pooledCount=pooled;
        }

        @Override
        public T input() {
            return t;
        }

        public void output(R ret) {
            rValues.offer(ret);
            int count=rValues.size();
            if(count==pooledCount){
                R sum=zero;
                for(int i=0;i<count;i++){
                    sum=aggFunc.call(sum,rValues.poll());
                }
                try{
                    Packet p1=Packet.wrapDataReplyPacket(sum,count,numOfIters);
                    context.write(p1,to);
                }catch (IOException e){
                    e.printStackTrace();
                }
            }
        }
    }

    private void thrRun(ClientIterativeRunnable<T,R> runnable, R zero, Function2<R,R,R> func){
        try(final IterationNetClient net= IterationNetClient.newClient(GlobalConfiguration.getProperty("spark.loop.type","udp"),
                localAddress,remoteAddress)){
            //noinspection Convert2Lambda
            net.waitForPacket(new IterationNetClient.OnPacketRead() {
                @Override
                public boolean onRead(IterationNetClient.VirtualNetContext ctx, Packet packet, SocketAddress from) {
                    IterationData<T,R> data=null;
                    int tmp;
                    LogEnv.getInstance().debug("Receiving %s from %s",packet,from);
                    switch (packet.type()){
                        case Packet.DataRequest:
                            LogEnv.getInstance().debug("Querying %d result",tmp=packet.numOfIterations());
                            data = new WrappedIterationData<>(packet.unwrapPacket(tClass),ctx,from,zero,func,tmp,pooledCount.get());
                            break;
                        case Packet.EchoRequest:
                            try {
                                int count=pooledCount.get(),count0=pooledThread.size()+1;
                                int echoSize=packet.echoSize();
                                ctx.write((echoSize == count && count==count0) ? Packet.EchoReplyOkPacket() : Packet.EchoReplyWaitingPacket(), from);
                                LogEnv.getInstance().debug("Sending %s(%d,%d,%d) to %s",(echoSize == count && count==count0)?"EchoOk":"EchoWaiting",
                                        echoSize,count,count0,from);
                            }catch (IOException e){
                                e.printStackTrace();
                            }
                            break;
                        case Packet.QuitRequest:
                            //noinspection unchecked
                            data=(IterationData<T, R>) EOFObject;
                            break;
                    }
                    if(data!=null){
                        for(BlockingQueue<IterationData<T,R>> queue:pooledThread.values()){
                            try{
                                queue.put(data);
                            }catch (InterruptedException ignored){}
                        }
                        if(runnable!=null && data!=EOFObject){
                            runnable.run(data);
                        }
                    }
                    return data!=EOFObject;
                }
            },0);
        }catch (IOException e){
            e.printStackTrace();
        }
    }

    private void waitingForData0(ClientIterativeRunnable<T,R> run, R zero, Function2<R,R,R> func){
        try {
            if (pooledCount.getAndIncrement() == 0) {
                thrRun(run, zero, func);
            } else {
                BlockingQueue<IterationData<T, R>> queue = new LinkedBlockingDeque<>(5);
                Thread current = Thread.currentThread();
                IterationData<T, R> data;
                try {
                    pooledThread.put(current, queue);
                    while ((data = queue.take()) != EOFObject) {
                        run.run(data);
                    }
                } catch (InterruptedException ignored) {
                }finally {
                    pooledThread.remove(current);
                }
            }
        }finally {
            if(pooledCount.decrementAndGet()==0){
                //clear
            }
        }
    }

    public void waitingForData(ClientIterativeRunnable<T,R> run, R zeroValue, Function2<R,R,R> func){
        waitingForData0(run,zeroValue,func);
    }

    private static final Map<String, IterationClient> clients=new HashMap<>();

    @SuppressWarnings("unchecked")
    public static <T,R> IterationClient<T,R> getOrCreateSingleInstance(String key, SocketAddress local, SocketAddress remote, Class<T> tcls){
        synchronized (clients){
            IterationClient<T,R> newClient=clients.get(key);
            if(newClient==null){
                newClient=new IterationClient<>();
                newClient.tClass=tcls;
                newClient.localAddress=local;
                newClient.remoteAddress=remote;
                clients.put(key,newClient);
                initLocalExecutor();
            }
            return newClient;
        }
    }

    public static Executor localExecutor;
    public static Map<Long, Executor.TaskRunner> runningTasks;
    @SuppressWarnings("unchecked")
    public static void initLocalExecutor(){
        NettyRpcEnv env=(NettyRpcEnv) SparkEnv.get().rpcEnv();
        try {
            Method func=NettyRpcEnv.class.getDeclaredMethod("dispatcher");
            func.setAccessible(true);
            Dispatcher dispatcher=(Dispatcher)func.invoke(env);
            func=Dispatcher.class.getDeclaredMethod("endpoints");
            func.setAccessible(true);
            Map<String,Object> endpoints=(Map<String,Object>)func.invoke(dispatcher);
            CoarseGrainedExecutorBackend ceb=(CoarseGrainedExecutorBackend) endpoints.get("Executor");
            Executor executor=ceb.executor();
            func=Executor.class.getDeclaredMethod("runningTasks");
            func.setAccessible(true);
            runningTasks=(Map<Long,Executor.TaskRunner>)func.invoke(executor);
            localExecutor=executor;
        } catch (NoSuchMethodException | IllegalAccessException | InvocationTargetException e) {
            e.printStackTrace();
        }
    }
}
